import nmap
import time
import sys
from Models import Target as Target
from Models import HostPort


class Port_Scan(object):
    def __init__(self):
        self.csv_write = []

    def scan(self, ip):

        nm = nmap.PortScanner()
        # try:
        # nm.scan(ip, arguments='1-65535')
        nm.scan(ip, arguments='22,80,135,139,443,445,593,901,1068,1935,3128,3306,4444,5800,5900,6129,6667,8080')
        print('[*]Host:%s' % ip)
        print(nm)

        for proto in nm[ip].all_protocols():
            print(nm[ip][proto].keys())
            lport = nm[ip][proto].keys()
            # lport.sort()
            for port in lport:
                self.csv_write.append([ip, port, nm[ip][proto][port]['state']])
                print('port : %s\tstate : %s' % (port, nm[ip][proto][port]['state']))
        # except:
        #     self.csv_write.append([ip, 'unknown', 'unknown'])
        #     print('port : unknown\tstate : unknown')

    def getPortInfo(self, ip):
        print('[+]Scanning...')
        startTime = time.time()

        self.scan(ip)

        endTime = time.time()
        print('------------------------------------------------------')
        print('[+]Scanning cost %ss ' % (endTime - startTime))
        print('[+]Done.')

        return self.csv_write


if __name__ == '__main__':
    host = sys.argv[1].strip()
    myScan = Port_Scan()
    openPortList = myScan.getPortInfo(host)
    HostPort.HostPort.insertMulti(openPortList)
